#include "Shader.h"
#include "Buffer.h"
#include "GL/glew.h"
#include "Texture.h"
#include "graphics/GLTextures.h"
#include "graphics/GLUniformBlocks.h"

#include <fstream>
#include <iostream>
#include <sstream>
#include <string.h>

#include <regex>
#include <unordered_map>

namespace Umikaze::Core {

class Uniform {
public:
  Uniform(int32_t location, uint32_t type) : _location(location), _type(type) {}

  void setBool(bool value) {
    checkUniformType(GL_BOOL);
    if (_cache.b != value) {
      glUniform1i(_location, (int)value);
      _cache.b = value;
    }
  }
  // ------------------------------------------------------------------------
  void setInt(int value) {
    checkUniformType(GL_INT);
    if (_cache.i != value) {
      glUniform1i(_location, value);
      _cache.i = value;
    }
  }
  // ------------------------------------------------------------------------
  void setFloat(float value) {
    checkUniformType(GL_FLOAT);
    if (_cache.f != value) {
      glUniform1f(_location, value);
      _cache.f = value;
    }
  }

  // ------------------------------------------------------------------------
  void setVec2(const glm::vec2 &value) {
    checkUniformType(GL_FLOAT_VEC2);
    if (_cache.v2 != value) {
      glUniform2fv(_location, 1, &value[0]);
      _cache.v2 = value;
    }
  }
  void setVec2(float x, float y) {
    checkUniformType(GL_FLOAT_VEC2);
    glm::vec2 temp(x, y);
    setVec2(temp);
  }
  // ------------------------------------------------------------------------
  void setVec3(const glm::vec3 &value) {
    checkUniformType(GL_FLOAT_VEC3);
    if (_cache.v3 != value) {
      glUniform3fv(_location, 1, &value[0]);
      _cache.v3 = value;
    }
  }
  void setVec3(float x, float y, float z) {
    checkUniformType(GL_FLOAT_VEC3);
    glm::vec3 temp(x, y, z);
    setVec3(temp);
  }
  // ------------------------------------------------------------------------
  void setVec4(const glm::vec4 &value) {
    checkUniformType(GL_FLOAT_VEC4);
    if (_cache.v4 != value) {
      glUniform4fv(_location, 1, &value[0]);
      _cache.v4 = value;
    }
  }
  void setVec4(float x, float y, float z, float w) {
    checkUniformType(GL_FLOAT_VEC4);
    glm::vec4 temp(x, y, z, w);
    setVec4(temp);
  }
  // ------------------------------------------------------------------------
  //   void setMat2( const glm::mat2 &mat) const{

  //   }
  //   //
  // ------------------------------------------------------------------------
  //   void setMat3( const glm::mat3 &mat) const;
  // ------------------------------------------------------------------------
  void setMat4(const glm::mat4 &mat) {
    checkUniformType(GL_FLOAT_MAT4);
    if (_cache.m4 != mat) {
      glUniformMatrix4fv(_location, 1, false, &mat[0][0]);
      _cache.m4 = mat;
    }
  }

  void setTexture(uint32_t textureLocation) const {
    if (_type != GL_SAMPLER_CUBE && _type != GL_SAMPLER_2D &&
        _type != GL_SAMPLER_2D_ARRAY) {
      std::cerr << "uniform类型不匹配" << std::endl;
    }
    glUniform1i(_location, textureLocation);
  }

private:
  inline void checkUniformType(GLenum type) const {
    if (type != _type) {
      std::cerr << "uniform类型不匹配" << std::endl;
    }
  };

  union Cache {
    bool b;
    int i;
    float f;
    glm::vec2 v2;
    glm::vec3 v3;
    glm::vec4 v4;
    glm::mat4 m4;
  } _cache;

  int32_t _location;
  uint32_t _type;
};

class UniformBlock {
public:
  UniformBlock(uint32_t index) : _index(index) {}

  void setUniformBlock(uint32_t programId, uint32_t location) {
    if (_cacheLocation != location) {
      _cacheLocation = location;
      glUniformBlockBinding(programId, _index, location);
    }
  }

  uint32_t _cacheLocation = 99999999;

  uint32_t _index;
};

class Attribute {
public:
  Attribute(int32_t attributeId, int32_t size, uint32_t type)
      : _attributeId(attributeId), _size(size), _type(type) {}

private:
  int32_t _attributeId;
  int32_t _size;
  uint32_t _type;
};

class IIShader {
public:
  IIShader(const std::string &vertCode, const std::string &fragCode,
           const std::string &geomCode = std::string()) {

    const char *vShaderCode = vertCode.c_str();
    const char *fShaderCode = fragCode.c_str();
    // 2. compile shaders
    unsigned int vertex, fragment;
    // vertex shader
    vertex = glCreateShader(GL_VERTEX_SHADER);
    glShaderSource(vertex, 1, &vShaderCode, NULL);
    glCompileShader(vertex);
    checkCompileErrors(vertex, "VERTEX");
    // fragment Shader
    fragment = glCreateShader(GL_FRAGMENT_SHADER);
    glShaderSource(fragment, 1, &fShaderCode, NULL);
    glCompileShader(fragment);
    checkCompileErrors(fragment, "FRAGMENT");
    // if geometry shader is given, compile geometry shader
    unsigned int geometry;
    if (!geomCode.empty()) {
      const char *gShaderCode = geomCode.c_str();
      geometry = glCreateShader(GL_GEOMETRY_SHADER);
      glShaderSource(geometry, 1, &gShaderCode, NULL);
      glCompileShader(geometry);
      checkCompileErrors(geometry, "GEOMETRY");
    }
    // shader Program
    uint32_t id = glCreateProgram();
    glAttachShader(id, vertex);
    glAttachShader(id, fragment);
    if (!geomCode.empty()) {
      glAttachShader(id, geometry);
    }
    glLinkProgram(id);
    checkCompileErrors(id, "PROGRAM");
    _programId = id;
    genUniforms();
  }

  ~IIShader() { glDeleteProgram(_programId); }

  void bind() {
    GlTextures::instance()->resetLocation();
    GlState::instance()->useProgram(_programId);
  }

  void setBool(const std::string &name, bool value) {
    if (auto iter = _uniforms.find(name); iter != _uniforms.end()) {
      iter->second->setBool(value);
    } else {
      uniformNameError(name);
    }
  }
  // ------------------------------------------------------------------------
  void setInt(const std::string &name, int value) const {
    if (auto iter = _uniforms.find(name); iter != _uniforms.end()) {
      iter->second->setInt(value);
    } else {
      uniformNameError(name);
    }
  }
  // ------------------------------------------------------------------------
  void setFloat(const std::string &name, float value) const {
    if (auto iter = _uniforms.find(name); iter != _uniforms.end()) {
      iter->second->setFloat(value);
    } else {
      uniformNameError(name);
    }
  }

  // ------------------------------------------------------------------------
  void setVec2(const std::string &name, const glm::vec2 &value) const {
    if (auto iter = _uniforms.find(name); iter != _uniforms.end()) {
      iter->second->setVec2(value);
    } else {
      uniformNameError(name);
    }
  }
  void setVec2(const std::string &name, float x, float y) const {
    if (auto iter = _uniforms.find(name); iter != _uniforms.end()) {
      iter->second->setVec2(x, y);
    } else {
      uniformNameError(name);
    }
  }
  // ------------------------------------------------------------------------
  void setVec3(const std::string &name, const glm::vec3 &value) const {
    if (auto iter = _uniforms.find(name); iter != _uniforms.end()) {
      iter->second->setVec3(value);
    } else {
      uniformNameError(name);
    }
  }
  void setVec3(const std::string &name, float x, float y, float z) const {
    if (auto iter = _uniforms.find(name); iter != _uniforms.end()) {
      iter->second->setVec3(x, y, z);
    } else {
      uniformNameError(name);
    }
  }
  // ------------------------------------------------------------------------
  void setVec4(const std::string &name, const glm::vec4 &value) const {
    if (auto iter = _uniforms.find(name); iter != _uniforms.end()) {
      iter->second->setVec4(value);
    } else {
      uniformNameError(name);
    }
  }
  void setVec4(const std::string &name, float x, float y, float z,
               float w) const {
    if (auto iter = _uniforms.find(name); iter != _uniforms.end()) {
      iter->second->setVec4(x, y, z, w);
    } else {
      uniformNameError(name);
    }
  }

  void setMat4(const std::string &name, const glm::mat4 &mat) const {
    if (auto iter = _uniforms.find(name); iter != _uniforms.end()) {
      iter->second->setMat4(mat);
    } else {
      uniformNameError(name);
    }
  }

  void setTexture(const std::string &name,
                  const std::shared_ptr<Texture> &texture) const {
    if (auto iter = _uniforms.find(name); iter != _uniforms.end()) {
      auto &glTextures = GlTextures::instance();
      uint32_t location = glTextures->getLocation();
      texture->bind(location);
      iter->second->setTexture(location);
    } else {
      uniformNameError(name);
    }
  }

  void setBuffer(const std::string &name,
                 const std::shared_ptr<Buffer> &buffer) const {
    if (auto iter = _uniformBlocks.find(name); iter != _uniformBlocks.end()) {
      auto ubo = buffer->getUbo();
      uint32_t location = GlUniformBocks::instance()->getLocation(ubo);
      iter->second->setUniformBlock(_programId, location);
    } else {
      uniformNameError(name);
    }
  }

private:
  void genUniforms() {
    int32_t attribCount = 0;
    glGetProgramiv(_programId, GL_ACTIVE_ATTRIBUTES, &attribCount);
    const int32_t bufSize = 32;
    int32_t length = -1;
    char name[bufSize];
    int32_t size = -1;
    GLenum type;
    for (int32_t i = 0; i < attribCount; ++i) {
      glGetActiveAttrib(_programId, static_cast<uint32_t>(i), bufSize, &length,
                        &size, &type, name);

      std::string trueName(name, length);
      int32_t attribLocation =
          glGetAttribLocation(_programId, trueName.c_str());
      _attribs.insert(
          {trueName, std::make_unique<Attribute>(attribLocation, size, type)});
      size = -1;
      length = -1;
      type = GLenum();
      memset(name, 0, bufSize);
    }

    int32_t uniformCount = 0;
    glGetProgramiv(_programId, GL_ACTIVE_UNIFORMS, &uniformCount);
    for (int32_t i = 0; i < uniformCount; ++i) {
      glGetActiveUniform(_programId, static_cast<uint32_t>(i), bufSize, &length,
                         &size, &type, name);

      std::string trueName;
      if (size > 1) {
        // TODO: 对于数组的情况下
        // nv会返回带 [0] 的名字
        // amd不会返回 [0] 的

        trueName = std::string(name, length - 3);

        for (int32_t i = 0; i < size; ++i) {
          std::string tempName = trueName + "[" + std::to_string(i) + "]";
          int32_t uniformLocation =
              glGetUniformLocation(_programId, tempName.c_str());
          _uniforms.insert(
              {tempName, std::make_unique<Uniform>(uniformLocation, type)});
        }

      } else {
        trueName = std::string(name, length);
        int32_t uniformLocation =
            glGetUniformLocation(_programId, trueName.c_str());

        _uniforms.insert(
            {trueName, std::make_unique<Uniform>(uniformLocation, type)});
      }

      size = -1;
      length = -1;
      type = GLenum();
      memset(name, 0, bufSize);
    }

    int32_t uniformBlockCount = 0;
    glGetProgramiv(_programId, GL_ACTIVE_UNIFORM_BLOCKS, &uniformBlockCount);
    for (int32_t i = 0; i < uniformBlockCount; ++i) {
      glGetActiveUniformBlockName(_programId, static_cast<uint32_t>(i), bufSize,
                                  &length, name);
      std::string trueName(name, length);

      uint32_t index = glGetUniformBlockIndex(_programId, trueName.c_str());
      _uniformBlocks.insert({trueName, std::make_unique<UniformBlock>(index)});
      length = -1;
      memset(name, 0, bufSize);
    }
  }

private:
  void uniformNameError(const std::string &name) const {
    // std::cerr << "Shader \"" << name << "\" 不包含uniform \"" << name << "\""
    //           << std::endl;
  }

  void checkCompileErrors(GLuint shader, std::string type) {
    GLint success;
    GLchar infoLog[1024];
    if (type != "PROGRAM") {
      glGetShaderiv(shader, GL_COMPILE_STATUS, &success);
      if (!success) {
        glGetShaderInfoLog(shader, 1024, NULL, infoLog);
        std::cout
            << "ERROR::SHADER_COMPILATION_ERROR of type: " << type << "\n"
            << infoLog
            << "\n -- --------------------------------------------------- -- "
            << std::endl;
      }
    } else {
      glGetProgramiv(shader, GL_LINK_STATUS, &success);
      if (!success) {
        glGetProgramInfoLog(shader, 1024, NULL, infoLog);
        std::cout
            << "ERROR::PROGRAM_LINKING_ERROR of type: " << type << "\n"
            << infoLog
            << "\n -- --------------------------------------------------- -- "
            << std::endl;
      }
    }
  }

  std::unordered_map<std::string, std::unique_ptr<Uniform>> _uniforms;
  std::unordered_map<std::string, std::unique_ptr<Attribute>> _attribs;
  std::unordered_map<std::string, std::unique_ptr<UniformBlock>> _uniformBlocks;
  uint32_t _programId;
};

class InternalShader {
public:
  InternalShader(const std::string &vertPath, const std::string &fragPath,
                 const std::string &geompath) {
    uint32_t offset = 0;
    _vertCode = load(vertPath, _vertConditions, offset);
    _fragCode = load(fragPath, _fragConditions, offset);
    if (!geompath.empty()) {
      _geomCode = load(geompath, _geomConditions, offset);
    }
  }

  std::unique_ptr<IIShader> &getShader(uint32_t version) {
    if (auto iter = _shaders.find(version); iter != _shaders.end()) {
      return iter->second;
    }

    std::stringstream vert;
    std::stringstream frag;

    vert << "#version 330 core\n";
    frag << "#version 330 core\n";
    addDefines(vert, _vertConditions, version);
    addDefines(frag, _fragConditions, version);
    vert << _vertCode;
    frag << _fragCode;

    uint32_t shaderId;
    if (!_geomCode.empty()) {
      std::stringstream geom;
      geom << "#version 330 core\n";
      addDefines(geom, _geomConditions, version);
      geom << _geomCode;
      _shaders.insert({version, std::make_unique<IIShader>(
                                    vert.str(), frag.str(), geom.str())});
      return _shaders.find(version)->second;
    } else {
      _shaders.insert(
          {version, std::make_unique<IIShader>(vert.str(), frag.str())});
      return _shaders.find(version)->second;
    }
  }

  uint32_t getFlag(const std::string &conditionName) {
    if (auto iter = _flags.find(conditionName); iter != _flags.end()) {
      return iter->second;
    }
    return 0;
  }

private:
  std::unordered_map<uint32_t, std::unique_ptr<IIShader>> _shaders;

private:
  void createDefine(const std::string &code,
                    std::unordered_map<uint32_t, std::string> &map,
                    uint32_t offset) {
    // TODO: 目前来说，只有condition相关的匹配，另外匹配了任意字符串也比
    // 最好是把限定一下宏的形式，不过正则写起来太麻烦，目前先这样，写的时候自己注意一下就好
    static const std::regex pattern(
        "\\s*(?:(?:#ifdef)|(?:#ifndef))\\s+(.*?)\n");
    std::smatch result;
    while (std::regex_search(code, result, pattern)) {
    }
  }

  void addDefines(std::stringstream &code,
                  const std::unordered_map<uint32_t, std::string> &map,
                  uint32_t version) {
    for (auto &iter : map) {
      if (version & iter.first) {
        code << "#define " << iter.second << "\n";
      }
    }
  }

  //每一个宏对应一个偏移量
  std::unordered_map<uint32_t, std::string> _vertConditions;
  std::unordered_map<uint32_t, std::string> _fragConditions;
  std::unordered_map<uint32_t, std::string> _geomConditions;
  std::unordered_map<std::string, uint32_t> _flags;

private:
  std::string getHeader(const std::string &headerName,
                        std::unordered_map<uint32_t, std::string> &map,
                        uint32_t &offset) {
    if (auto iter = _headers.find(headerName); iter != _headers.end()) {
      return iter->second;
    }
    auto header = load(headerName, map, offset);
    _headers.insert({headerName, header});
    return header;
  }

  std::string load(const std::string &name,
                   std::unordered_map<uint32_t, std::string> &map,
                   uint32_t &offset) {
    // TODO:
    // 默认shader文件有正常的换行格式，另外现在匹配了任意字符穿，其实最后有所限制
    // 不过写正则很麻烦，暂时先这样，能用就行
    static const std::regex includePattern(
        "\\s*#include\\s+(?:(?:<([A-Z_a-z][A-z_a-z0-9\\./"
        "]+)>)|(?:\"([A-Z_a-z][A-z_a-z0-9\\./]+)\"))");

    // TODO:
    // 同上，默认了换行，也同样没有限制规则
    static const std::regex defPattern(
        "(?:#ifn?def\\s+([A-Z_a-z][A-z_a-z0-9]+))|(?:defined\\(([A-Z_a-z][A-z_"
        "a-z0-9]+)\\))");
    std::stringstream ss;
    try {
      uint32_t err = 0;
      std::ifstream codeFile(name);

      if (!codeFile.is_open()) {
        throw "Cant open file";
      }

      std::string line;
      while (std::getline(codeFile, line)) {
        if (line.empty()) {
          continue;
        }

        //检验是不是 #include语句
        if (std::smatch results;
            std::regex_search(line, results, includePattern)) {
          std::string headerName;
          if (results[1].length() != 0) {
            headerName = results[1];
          } else if (results[2].length() != 0) {
            headerName = results[2];
          }
          // 这个函数会重新调回到load，会自动把头里的宏添加到map
          ss << getHeader(headerName, map, offset);
        } else {
          // 需要检验这一句是不是包含宏
          ss << line << "\n";
          std::smatch res;
          while (std::regex_search(line, res, defPattern)) {
            if (res[1].length() != 0 || res[2].length() != 0) {
              std::string macro = res[1].length() == 0 ? res[2] : res[1];
              if (auto iter = _flags.find(macro); iter != _flags.end()) {
                map.insert({iter->second, macro});
              } else {
                map.insert({1 << offset, macro});
                _flags.insert({macro, 1 << offset++});
              }
              line = res.suffix().str();
            }
          }
        }
      }

      if (codeFile.bad()) {
        throw "Read fail";
      }

      return ss.str();
    } catch (std::ifstream::failure &e) {
      auto what = e.what();
      std::cerr << "ERROR::SHADER::FILE_NOT_SUCCESFULLY_READ" << std::endl;
      return std::string();
    }
  }

  std::string _vertCode;
  std::string _fragCode;
  std::string _geomCode;
  static std::unordered_map<std::string, std::string> _headers;

private:
};

std::unordered_map<std::string, std::string> InternalShader::_headers;

Shader::Shader(const std::string &vertPath, const std::string &fragPath,
               const std::string &geomPath)
    : _internalShader(
          std::make_unique<InternalShader>(vertPath, fragPath, geomPath)) {}

void Shader::bind() {

  auto &iishader = _internalShader->getShader(_version);
  iishader->bind();

  for (auto iter : _bools) {
    iishader->setBool(iter.first, iter.second);
  }

  for (auto iter : _ints) {
    iishader->setInt(iter.first, iter.second);
  }

  for (auto iter : _floats) {
    iishader->setFloat(iter.first, iter.second);
  }

  for (auto iter : _vec2s) {
    iishader->setVec2(iter.first, iter.second);
  }

  for (auto iter : _vec3s) {
    iishader->setVec3(iter.first, iter.second);
  }

  for (auto iter : _vec4s) {
    iishader->setVec4(iter.first, iter.second);
  }

  for (auto iter : _mat4s) {
    iishader->setMat4(iter.first, iter.second);
  }

  for (auto iter : _textures) {
    if (!iter.second.expired()) {
      iishader->setTexture(iter.first, iter.second.lock());
    }
  }

  for (auto iter : _blocks) {
    if (!iter.second.expired()) {
      iishader->setBuffer(iter.first, iter.second.lock());
    }
  }
}

void Shader::setBool(const std::string &name, bool value) {
  if (auto iter = _bools.find(name); iter != _bools.end()) {
    iter->second = value;
  } else {
    _bools.insert({name, value});
  }
}
// ------------------------------------------------------------------------
void Shader::setInt(const std::string &name, int value) {
  if (auto iter = _ints.find(name); iter != _ints.end()) {
    iter->second = value;
  } else {
    _ints.insert({name, value});
  }
}
// ------------------------------------------------------------------------
void Shader::setFloat(const std::string &name, float value) {
  if (auto iter = _floats.find(name); iter != _floats.end()) {
    iter->second = value;
  } else {
    _floats.insert({name, value});
  }
}

// ------------------------------------------------------------------------
void Shader::setVec2(const std::string &name, const glm::vec2 &value) {
  if (auto iter = _vec2s.find(name); iter != _vec2s.end()) {
    iter->second = value;
  } else {
    _vec2s.insert({name, value});
  }
}
void Shader::setVec2(const std::string &name, float x, float y) {
  setVec2(name, {x, y});
}
// ------------------------------------------------------------------------
void Shader::setVec3(const std::string &name, const glm::vec3 &value) {
  if (auto iter = _vec3s.find(name); iter != _vec3s.end()) {
    iter->second = value;
  } else {
    _vec3s.insert({name, value});
  }
}
void Shader::setVec3(const std::string &name, float x, float y, float z) {
  setVec3(name, {x, y, z});
}
// ------------------------------------------------------------------------
void Shader::setVec4(const std::string &name, const glm::vec4 &value) {
  if (auto iter = _vec4s.find(name); iter != _vec4s.end()) {
    iter->second = value;
  } else {
    _vec4s.insert({name, value});
  }
}
void Shader::setVec4(const std::string &name, float x, float y, float z,
                     float w) {
  setVec4(name, {x, y, z, w});
}

void Shader::setMat4(const std::string &name, const glm::mat4 &value) {
  if (auto iter = _mat4s.find(name); iter != _mat4s.end()) {
    iter->second = value;
  } else {
    _mat4s.insert({name, value});
  }
}

void Shader::setFloatArray(const std::string &name,
                           const std::vector<float> &number) {
  for (uint32_t i = 0; i < number.size(); ++i) {
    std::string trueName = name + "[" + std::to_string(i) + "]";
    if (auto iter = _floats.find(trueName); iter != _floats.end()) {
      iter->second = number[i];
    } else {
      _floats.insert({trueName, number[i]});
    }
  }
}

void Shader::setTexture(const std::string &name,
                        const std::shared_ptr<Texture> &texture) {
  if (auto iter = _textures.find(name); iter != _textures.end()) {
    iter->second = texture;
  } else {
    _textures.insert({name, texture});
  }
}

void Shader::setBuffer(const std::string &name,
                       const std::shared_ptr<Buffer> &buffer) {

  if (auto iter = _blocks.find(name); iter != _blocks.end()) {
    iter->second = buffer;
  } else {
    _blocks.insert({name, buffer});
  }
}

void Shader::setCondition(const std::string &name, bool state) {
  auto flag = _internalShader->getFlag(name);
  if (!flag) {
    std::cerr << "没有这个condition" << std::endl;
    return;
  }
  if (state) {
    _version |= flag;
  } else {
    _version &= ~flag;
  }
}

Shader::~Shader() = default;

} // namespace Umikaze::Core
